'''compute_drawdowns.py
Description: Script used to compute security drawdowns
'''
import pandas as pd
import numpy as np
import argparse

# drawdown compute function
def compute_drawdowns(input_file, id_field, date_field, price_field, out_file):
    
    raw_data = pd.read_stata(input_file)
    raw_data = raw_data.sort_values([id_field, date_field])

    if raw_data[price_field].min() < 0:
        raw_data[price_field] = raw_data[price_field] - raw_data[price_field].min()
    
    raw_data['cum_max_price'] = raw_data.groupby([id_field])[price_field].cummax()
    raw_data['sub_drawdown'] = 1 - raw_data[price_field] / raw_data.cum_max_price
    
    drawdowns = pd.DataFrame(
            raw_data.groupby([id_field])['sub_drawdown'].max()
        ).reset_index().rename(columns={'sub_drawdown': 'py_drawdown'})
    drawdowns.to_stata(out_file)
    
if __name__ == "__main__":
    
    # set up parser
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_file', type=str)
    parser.add_argument('--id_field', type=str)
    parser.add_argument('--date_field', type=str)
    parser.add_argument('--price_field', type=str)
    parser.add_argument('--out_file', type=str)

    # run parsing
    args = parser.parse_args()
    print("=== Computing drawdowns ===")
    print("Input file: ", args.input_file)
    print("ID field: ", args.id_field)
    print("Date field: ", args.date_field)
    print("Price field: ", args.price_field)
    print("Output file: ", args.out_file)

    # run function
    compute_drawdowns(
        input_file = args.input_file,
        id_field = args.id_field,
        date_field = args.date_field,
        price_field = args.price_field,
        out_file = args.out_file
    )
